using Crux.Core.LangChain.Abstractions.Databases;
using Crux.Core.LangChain.Abstractions.Databases.Metadatas;
using MongoDB.Bson.Serialization;
using MongoDB.Driver;

namespace Crux.Core.LangChain.Databases.Mongo;
/// <summary>
/// 
/// </summary>
public class MongoRagCollection(
    IMongoDatabase database,
    string indexName,
    string name = MongoRagCollection.DefaultCollection) : IVectorCollection
{
    private const string DefaultCollection = "langchain";
    private readonly IMongoCollection<Vector> _collection = database.GetCollection<Vector>(name);
    /// <summary>
    /// 
    /// </summary>
    /// <param name="items"></param>
    /// <param name="cancellationToken"></param>
    /// <returns></returns>
    public async Task<IReadOnlyCollection<string>> AddAsync(IReadOnlyCollection<Vector> items, CancellationToken cancellationToken = default)
    {
        await _collection.InsertManyAsync(items, cancellationToken: cancellationToken).ConfigureAwait(false);
        return items.Select(i => i.Id).ToArray();
    }
    /// <summary>
    /// 
    /// </summary>
    /// <param name="ids"></param>
    /// <param name="cancellationToken"></param>
    /// <returns></returns>
    public async Task<bool> DeleteAsync(IEnumerable<string> ids, CancellationToken cancellationToken = default)
    {
        var filter = Builders<Vector>.Filter.In(i => i.Id, ids);
        var result = await _collection.DeleteManyAsync(filter, cancellationToken).ConfigureAwait(false);

        return result.IsAcknowledged;
    }
    /// <summary>
    /// 
    /// </summary>
    /// <param name="id"></param>
    /// <param name="cancellationToken"></param>
    /// <returns></returns>
    public async Task<Vector?> GetAsync(string id, CancellationToken cancellationToken = default)
    {
        var filter = Builders<Vector>.Filter.Eq(i => i.Id, id);
        var result = await _collection.FindAsync(filter, cancellationToken: cancellationToken).ConfigureAwait(false);
        return result.FirstOrDefault(cancellationToken: cancellationToken);
    }
    /// <summary>
    /// 
    /// </summary>
    /// <param name="cancellationToken"></param>
    /// <returns></returns>
    /// <exception cref="NotImplementedException"></exception>
    public async Task<bool> IsEmptyAsync(CancellationToken cancellationToken = default)
    {
        return await _collection.EstimatedDocumentCountAsync(cancellationToken: cancellationToken).ConfigureAwait(false) == 0;
    }
    /// <summary>
    /// 
    /// </summary>
    /// <param name="request"></param>
    /// <param name="settings"></param>
    /// <param name="cancellationToken"></param>
    /// <returns></returns>
    /// <exception cref="ArgumentNullException"></exception>
    public async Task<VectorSearchResponse> SearchAsync(VectorSearchRequest request, VectorSearchSettings? settings = null, CancellationToken cancellationToken = default)
    {
        request = request ?? throw new ArgumentNullException(nameof(request));
        settings ??= new VectorSearchSettings();

        var options = new VectorSearchOptions<Vector>()
        {
            IndexName = indexName,
        };

        var projectionDefinition = Builders<Vector>.Projection
                         .Exclude(a => a.Distance)
                         .Meta("Distance", "vectorSearchScore")
                         .Include(e => e.Text);

        var results = await _collection.Aggregate()
                   .VectorSearch(z => z.Embedding, request.Embeddings.First(), settings.NumberOfResults, options)
                   .Project<Vector>(projectionDefinition)
                   .ToListAsync(cancellationToken)
                   .ConfigureAwait(false);

        return new VectorSearchResponse
        {
            Items = results.ToArray(),
        };
    }
    /// <summary>
    /// 
    /// </summary>
    /// <param name="filters"></param>
    /// <param name="cancellationToken"></param>
    /// <returns></returns>
    /// <exception cref="ArgumentNullException"></exception>
    /// <exception cref="ArgumentException"></exception>
    public async Task<IReadOnlyList<Vector>> SearchByMetadata(Dictionary<string, object> filters, CancellationToken cancellationToken = default)
    {
        filters = filters ?? throw new ArgumentNullException(nameof(filters));

        var builder = Builders<Vector>.Filter;
        var filterDefinitions = new List<FilterDefinition<Vector>>();

        foreach (var kvp in filters)
        {
            if (kvp.Value == null)
            {
                throw new ArgumentException($"Metadata value for key '{kvp.Key}' cannot be null", nameof(filters));
            }

            var filter = builder.Eq($"Metadata.{kvp.Key}", kvp.Value);
            filterDefinitions.Add(filter);
        }

        var combinedFilter = builder.And(filterDefinitions);

        var results = await _collection
            .Find(combinedFilter)
            .ToListAsync(cancellationToken)
            .ConfigureAwait(false);

        return results;
    }
}

